# 文件名: modules/trap.py
import pygame
import random
from .utils import get_random_grid_position, draw_grid_cell

'''陷阱管理'''

def place_traps(cfg, count, snake_body, existing_traps=None, food_pos=None, poison_pos=None):
    """
    放置指定数量的陷阱（网格坐标）。
    确保不与蛇初始位置、现有陷阱、食物、毒药重叠。
    """
    traps = [] if existing_traps is None else list(existing_traps) # Start with existing traps if any
    occupied = set(map(tuple, snake_body)) # Use set for faster lookup
    if food_pos: occupied.add(tuple(food_pos))
    if poison_pos: occupied.add(tuple(poison_pos))
    if existing_traps: occupied.update(map(tuple, existing_traps))

    attempts = 0
    max_attempts = cfg.GRID_WIDTH * cfg.GRID_HEIGHT * 2 # Limit attempts

    while len(traps) < count and attempts < max_attempts:
        position = get_random_grid_position(cfg)
        pos_tuple = tuple(position)
        # Ensure trap is not in an already occupied space
        if pos_tuple not in occupied:
            traps.append(position)
            occupied.add(pos_tuple) # Add new trap to occupied set
        attempts += 1

    if len(traps) < count:
        print(f"Warning: Could only place {len(traps)} out of {count} requested traps.")

    return traps # Return list of [x, y] coordinates

def draw_traps(screen, traps, cfg):
    """绘制所有陷阱"""
    for trap_pos in traps:
        # Corrected the call to draw_grid_cell: pass cfg, not cfg.cfg
        draw_grid_cell(screen, trap_pos, cfg.TRAP_COLOR, cfg)